# test.py

import torch
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from dataset import CustomDataset
from model import CustomNet
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import numpy as np
from torchvision import transforms

def evaluate_model(dataloader, model, device):  # 重命名函数
    """测试模型并输出评估指标"""
    model.eval()
    size = len(dataloader.dataset)
    correct = 0
    all_labels = []
    all_preds = []

    with torch.no_grad():
        for data in dataloader:
            inputs = data['image'].to(device)
            labels = data['label'].to(device)

            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)

            correct += (predicted == labels).sum().item()
            all_labels.extend(labels.cpu().numpy())
            all_preds.extend(predicted.cpu().numpy())

    # 计算准确率
    accuracy = 100. * correct / size
    print(f'Accuracy: {accuracy:.2f}%')

    # 打印分类报告
    print("\nClassification Report:")
    print(classification_report(all_labels, all_preds))

    # 绘制混淆矩阵
    cm = confusion_matrix(all_labels, all_preds)
    plt.figure(figsize=(10, 8))
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
    plt.title('Confusion matrix')
    plt.colorbar()

    classes = [str(i) for i in range(len(np.unique(all_labels)))]
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes)

    # 在混淆矩阵上标注数字
    thresh = cm.max() / 2.
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            plt.text(j, i, format(cm[i, j], 'd'),
                     horizontalalignment="center",
                     color="white" if cm[i, j] > thresh else "black")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    plt.savefig('confusion_matrix.png')
    plt.show()

    return accuracy


if __name__ == "__main__":
    # 数据预处理（与训练时保持一致）
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 加载测试数据
    test_dataset = CustomDataset('./images/test.txt', './images/test', lambda: transform)
    test_dataloader = DataLoader(test_dataset, batch_size=16)

    # 加载模型
    model = CustomNet()
    model.load_state_dict(torch.load('./models/best_model.pth', map_location=torch.device('cpu')))
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # 运行测试
    evaluate_model(test_dataloader, model, device)  # 调用重命名后的函数